{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quickstart: $J/\\Psi \\rightarrow \\gamma \\pi^0 \\pi^0$ decay\n",
    "\n",
    "In this quickstart example it is shown how to use ComPWA via the python interface. The workflow is:\n",
    "\n",
    "1. Create a model for the decay.\n",
    "2. Generate a Monte Carlo data sample (hit & miss) using this model. \n",
    "3. Perform a fit on the data sample using the Minuit2 interface.\n",
    "4. Visualize the data and the fit result!\n",
    "\n",
    "Let's go!\n",
    "\n",
    "First we `import` the necessary expert system module parts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.expertsystem.ui.system_control import (\n",
    "    StateTransitionManager, InteractionTypes)\n",
    "from pycompwa.expertsystem.amplitude.helicitydecay import (\n",
    "    HelicityAmplitudeGeneratorXML)\n",
    "\n",
    "# just a little function to print the intermediate states\n",
    "def print_intermediate_states(solutions):\n",
    "    from pycompwa.expertsystem.topology.graph import (\n",
    "        get_intermediate_state_edges)\n",
    "    print(\"intermediate states:\")\n",
    "    intermediate_states = set()\n",
    "    for g in solutions:\n",
    "        edge_id = get_intermediate_state_edges(g)[0]\n",
    "        intermediate_states.add(g.edge_props[edge_id]['@Name'])\n",
    "    print(intermediate_states)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Creating the Decay Model\n",
    "\n",
    "#### 1.1. Define problem set\n",
    "\n",
    "First we define the boundary conditions of our physics problem, such as\n",
    "\n",
    "- initial state\n",
    "- final state\n",
    "- formalism type\n",
    "- ...\n",
    "\n",
    "Pass all of that information to the `StateTransitionManager`, which is the main user interface class of the ComPWA expert system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_state = [(\"J/psi\", [-1, 1])]\n",
    "final_state = [(\"gamma\", [-1, 1]), (\"pi0\", [0]), (\"pi0\", [0])]\n",
    "\n",
    "tbd_manager = StateTransitionManager(initial_state, final_state,\n",
    "                                     formalism_type='helicity',\n",
    "                                     topology_building='isobar')"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    ".. note::\n",
    "   The ``StateTransitionManager`` (STM) is the main user interface class of the\n",
    "   ComPWA expert system. The boundary conditions of your physics problem are \n",
    "   defined here, such as the initial state, final state, formalism type, ...\n",
    "\n",
    "   * ``prepare_graphs()`` of the STM creates all topology graphs, here using\n",
    "     the isobar model (two-body decays). Also it initializes the graphs with\n",
    "     the initial and final state and the a set of conservation laws at each\n",
    "     interaction node.\n",
    "\n",
    "   * By default all three (strong, EM, weak) interaction types are used in the\n",
    "     preparation stage. However it is also possible to globally choose the\n",
    "     allowed interaction types via ``set_allowed_interaction_types()``.\n",
    "\n",
    "   After the preparation step, you can modifiy the settings returned by\n",
    "   ``prepare_graphs()`` to your liking. Since this output is quite a lot of\n",
    "   information, the expertsystem ui is supposed to aid in the configuration\n",
    "   (especially the STM).\n",
    "   \n",
    "   * A subset of particles that are allow as intermediate states can also be\n",
    "     specified in the STM. Either in the ``init()`` of the STM or setting the\n",
    "     instance attribute ``allowed_intermediate_particles``."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.2. Preparation\n",
    "Create all topology graphs using the isobar model (two-body decays).\n",
    "\n",
    "Also initialize the graphs with the initial and final state. Remember that each interaction node defines their own set of conservation laws. The `StateTransitionManager` (STM) defines three interaction types:\n",
    "\n",
    "| Interaction | Strength |\n",
    "| --- | --- |\n",
    "| strong | $60$ |\n",
    "| electromagnetic (EM) | $1$ |\n",
    "| weak | $10^{-4}$ |\n",
    "\n",
    "Be default all three are used in the preparation stage. `prepare_graphs()` of the STM generates graphs with all possible combinations of interaction nodes. An overall interaction strength is assigned to each graph, and they are grouped according to this strength."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.3. Finding solutions\n",
    "If you are happy with the automatic settings generated by the StateTransitionManager, just go directly to the solving!"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    ".. note::\n",
    "   This step takes about 30 sec on a Intel(R) Core(TM) i7-6820HQ CPU @ 2.70GHz running multi-threaded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(solutions, violated_rules) = tbd_manager.find_solutions(\n",
    "        graph_interaction_settings_groups)\n",
    "\n",
    "print(\"found \" + str(len(solutions)) + \" solutions!\")\n",
    "print_intermediate_states(solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok, now we have a lot of solutions that are actually heavily supressed (involve two weak decays). In general you can modify the dictionary return by `prepare_graphs()` directly.\n",
    "\n",
    "The STM also comes with a functionality to globally choose the allowed interaction types (`set_allowed_interaction_types()`). Go ahead and **disable** the **EM** and **weak** interaction!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tbd_manager.set_allowed_interaction_types(\n",
    "    [InteractionTypes.Strong])\n",
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()\n",
    "(solutions, violated_rules) = tbd_manager.find_solutions(\n",
    "        graph_interaction_settings_groups)\n",
    "print(\"found \" + str(len(solutions)) + \" solutions!\")\n",
    "print_intermediate_states(solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Huh, what happened here? Actually, since a $\\gamma$ **particle** appears in one of the interaction nodes, the expert system knows that this node must involve **EM interaction**! Because the node can be an effective interaction, the weak interaction cannot be excluded, as it contains only a subset of conservation laws.\n",
    "\n",
    "Since only the strong interaction was supposed to be used, this gives a warning and automatically corrects the mistake.\n",
    "\n",
    "Once the EM interaction is included, this warning disappears. However be aware that the EM interaction is now available globally. Hence, now there might be solutions in which both nodes are electromagnetic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tbd_manager.set_allowed_interaction_types(\n",
    "    [InteractionTypes.Strong, InteractionTypes.EM])\n",
    "graph_interaction_settings_groups = tbd_manager.prepare_graphs()\n",
    "(solutions, violated_rules) = tbd_manager.find_solutions(\n",
    "        graph_interaction_settings_groups)\n",
    "\n",
    "print(\"found \" + str(len(solutions)) + \" solutions!\")\n",
    "print_intermediate_states(solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! Now we selected the solutions that are contribution the strongest. However, be aware that there are more effects that can suppress certain decays. For example branching ratios. In this example $J/\\Psi$ can decay into $\\pi^0 + \\rho^0$ or $\\pi^0 + \\omega$.\n",
    "\n",
    "| decay | branching ratio |\n",
    "| --- | --- |\n",
    "| $$\\omega \\rightarrow \\gamma+\\pi^0$$ | 0.0828 |\n",
    "| $$\\rho^0 \\rightarrow \\gamma+\\pi^0$$ | 0.0006 |\n",
    "\n",
    "Unfortunately the $\\rho^0$ decays mainly into $\\pi+\\pi$, not $\\gamma+\\pi^0$. Hence it is suppressed. This information is currently not known to the expert system.\n",
    "But it is possible to hand the expert system a list of allowed intermediate states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# particles are found by name comparison; so i.e. f2 will find all f2's and f all f's independent of their spin\n",
    "tbd_manager.allowed_intermediate_particles = ['f']\n",
    "\n",
    "(solutions, violated_rules) = tbd_manager.find_solutions(\n",
    "        graph_interaction_settings_groups)\n",
    "\n",
    "print(\"found \" + str(len(solutions)) + \" solutions!\")\n",
    "print_intermediate_states(solutions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have selected all amplitudes that involve **f** states. The appearing warnings notify the user, that the solution space is only partial. For certain lines in the graph, no suitable particle was found (since only f state were allowed).\n",
    "\n",
    "At this point we are all set to generate some data using this amplitude model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xml_generator = HelicityAmplitudeGeneratorXML()\n",
    "xml_generator.generate(solutions)\n",
    "xml_generator.write_to_file('model.xml')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Creating a data sample\n",
    "\n",
    "In this section we will use the amplitude model created above to generate a data sample via hit & miss Monte Carlo.\n",
    "\n",
    "Using this amplitude model in the c++ side of ComPWA is simple. The `create_intensity_and_kinematics()` helper function builds an Intensity and Kinematics object, as specified in the xml file. These can be used to generate our data sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycompwa.ui as pwa\n",
    "\n",
    "intensity, kinematics = pwa.create_intensity_and_kinematics('model.xml')\n",
    "\n",
    "generator = pwa.RootGenerator(kinematics.get_particle_state_transition_kinematics_info())\n",
    "\n",
    "random_generator = pwa.RootUniformRealGenerator(12345)\n",
    "\n",
    "datasample = pwa.generate(5000, kinematics, generator, intensity, random_generator)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    ".. note::\n",
    "   ``pycompwa.ui`` is the python interface to ComPWA's c++ modules. Read more\n",
    "   about this :ref:`here <python-ui>`.\n",
    "\n",
    "   Three important pieces for evaluating an intensity are:\n",
    "\n",
    "   * The **intensity** itself. It was generated previously and stored within\n",
    "     the xml model file.\n",
    "\n",
    "   * A **kinematics** instance. It handles the calculation of the kinematic\n",
    "     variables that are required for the evaluation of the intensity!\n",
    "     For example in the helicity formalism: :math:`(s,\\theta,\\phi)`.\n",
    "   \n",
    "   * **Data samples**. For mere visualization of the intensity a phase space\n",
    "     sample is sufficient. It is mandatory for the normalization of the\n",
    "     intensity. However when performing fits an additional data sample, to\n",
    "     which the intensity will be compared to, has to be specified."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Fitting\n",
    "\n",
    "An **Intensity** object behaves just like a mathematical function that takes set of data as an argument and returns a list of intensities (real numbers). To perform a fit create an estimator of your choice, which generally needs\n",
    "- an `Intensity` instance\n",
    "- a `DataSet` (to which the intensity is fitted)\n",
    "- optionally a phase space `DataSet` (which is used to normalize the Intensity)\n",
    "\n",
    "A phase space sample can be generated via the `generate_phsp()` function. The data samples can be converted to `DataSet` using the `convert_events_to_dataset()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phspsample = pwa.generate_phsp(100000, generator, random_generator)\n",
    "\n",
    "dataset = pwa.convert_events_to_dataset(datasample, kinematics)\n",
    "phspdataset = pwa.convert_events_to_dataset(phspsample, kinematics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now an estimator can be created. Remember that it indicates how well a set of model parameters describes a given data set best. In this example an unbinned log likelihood estimator is used. It performs the calculations with the FunctionTree backend. In this example the intensity is not normalized. Therefore the additional phase space dataset `phspdataset` is supplied, so that the estimator handles the normalization for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator, initial_parameters = pwa.create_unbinned_log_likelihood_function_tree_estimator(intensity, \n",
    "                                                                                           dataset,\n",
    "                                                                                           phspdataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that you not only receive a estimator object, but also a list of fit parameters (`FitParameterList`). This list of fit parameters are the settings which initialize the optimizer later on. They contain the following info:\n",
    "- the initial values of the parameters\n",
    "- fix parameters\n",
    "- define boundaries\n",
    "- define errors, which can give certain optimizers hints on the stepsize\n",
    "\n",
    "These fit parameters are initialized with the values stated in the xml file or default values. But can be changed easily, like fixing certain parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "pprint(initial_parameters)\n",
    "print(\"this parameter is initially not fixed: \", initial_parameters[1])\n",
    "initial_parameters[1].is_fixed = True\n",
    "print(\"and now it is fixed: \", initial_parameters[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make the fit a bit more interesting, we modify one of the parameters to a different initial value then the true value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(initial_parameters[6])\n",
    "initial_parameters[6].value = 2.0\n",
    "print(\"after:\",initial_parameters[6])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's time to start up a set up a fit, which is quite simple. Just create an optimizer instance of your choice, here Minuit2 (`MinuitIF`), and call the `optimize()` function to start the fitting process."
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    ".. note::\n",
    "   The fit (with 15 free parameters) on a Intel(R) Core(TM) i7-6820HQ CPU @ 2.70GHz takes about 10 sec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minuitif = pwa.MinuitIF()\n",
    "result = minuitif.optimize(estimator, initial_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.fit_duration_in_seconds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check if the our fit parameters are \"close to\" the true values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"this should be close to 1.0 again: \", \n",
    "      result.final_parameters[6].value, \n",
    "      \" +- \", \n",
    "      result.final_parameters[6].error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Important**: Note that the intensity instance still needs to be notified about this optimal set of parameters. They can be applied with the `updateParametersFrom()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensity.updateParametersFrom(result.final_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Visualization\n",
    "Let's go ahead and make a Dalitz plot of the data sample and our fit result. ComPWA's pycompwa package ships with a little plotting module to help you generate some common plots using matplotlib. Before we can utilize this, the data has to be extended.\n",
    "\n",
    "Since our model only includes one specific decay topology the `HelicityKinematics` class only generates the kinematic variables needed to evaluate the Intensity. In order to make a Dalitz plot, also the kinematic variables from the other subsystems are needed. They can be created with `create_all_subsystems()`. However the conversion from the event samples have to be performed again!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kinematics.create_all_subsystems()\n",
    "dataset = pwa.convert_events_to_dataset(datasample, kinematics)\n",
    "phspdataset = pwa.convert_events_to_dataset(phspsample, kinematics)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    ".. tip::\n",
    "   \n",
    "   ComPWA ships with a little python plotting module (``plotting``), which for example helps you to read in ROOT TTree's and create common plots (e.g. angular distributions, Dalitz plots). It uses matplotlib as a backend. You can either hand it data files, or feed it directly with data.\n",
    "\n",
    "   Please use it, instead of creating your own visualization scripts!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are two variants to create plots from data:\n",
    "\n",
    "#### Variant 1: From a ROOT file\n",
    "\n",
    "The ROOT file can either be generated directly via c++ code (`RootPlotData`), or you can use the `create_rootplotdata()` helper function of the python UI. So let's create a file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwa.create_rootplotdata(\"rootplot.root\", kinematics, data_sample=dataset, \n",
    "                        phsp_sample=phspdataset, intensity=intensity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the `rootplotdatareader` submodule of `plotting` can be used to read in the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.plotting.rootplotdatareader import open_compwa_plot_data\n",
    "\n",
    "plot_data = open_compwa_plot_data(\"rootplot.root\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The returned data structure can directly be fed to `make_dalitz_plots()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pycompwa.plotting import make_dalitz_plots\n",
    "make_dalitz_plots(plot_data, [\"mSq_3_4_vs_2\", \"mSq_2_4_vs_3\"], bins=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Variant 2: By direct data access\n",
    "\n",
    "The ComPWA python interface has two helper functions `create_data_array()` `create_fitresult_array()` that allow the datapoins to be exported as two-dimensional lists. These can be used to create a numpy record arrays. Then the **plotting** module is able to create common plots like the Dalitz plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Plotting\n",
    "from pycompwa.plotting import (\n",
    "    make_dalitz_plots, PlotData, create_nprecord,\n",
    "    make_difference_plot_2d\n",
    ")\n",
    "\n",
    "# use the direct data point access\n",
    "var_names, dataarray = pwa.create_data_array(dataset)\n",
    "data_record = create_nprecord(var_names, dataarray)\n",
    "\n",
    "fitres_var_names, fitres_dataarray = pwa.create_fitresult_array(intensity, phspdataset)\n",
    "fitresult_record = create_nprecord(fitres_var_names, fitres_dataarray)\n",
    "\n",
    "plot_data = PlotData(data_record=data_record, fit_result_data_record=fitresult_record)\n",
    "#plot_data.particle_id_to_name_mapping = data_points.get_finalstate_id_to_name_mapping()\n",
    "\n",
    "# plot a 2d histogram\n",
    "make_dalitz_plots(plot_data, [\"mSq_3_4_vs_2\", \"mSq_2_4_vs_3\"], bins=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make the differences more clearly visible we can create a 2d difference plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_difference_plot_2d(plot_data, [\"mSq_3_4_vs_2\", \"mSq_2_4_vs_3\"], bins=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it. You can check some of the other examples to learn about more detailed features of ComPWA.\n",
    "\n",
    "And give us feedback or contribute ;)!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
